#!/usr/bin/python
# -*-coding:utf-8-*-
"""Utilities that are required by gplearn.

Most of these functions are slightly modified versions of some key utility
functions from scikit-learn that gplearn depends upon. They reside here in
order to maintain compatibility across different versions of scikit-learn.

"""

import numbers

import numpy as np
from joblib import cpu_count

from .base.GPFunctionsV1 import *
from .base.GPFunctionsV2 import *
from .base.GPFunctionsV3 import *


def check_random_state(seed):
    """Turn seed into a np.random.RandomState instance

    Parameters
    ----------
    seed : None | int | instance of RandomState
        If seed is None, return the RandomState singleton used by np.random.
        If seed is an int, return a new RandomState instance seeded with seed.
        If seed is already a RandomState instance, return it.
        Otherwise raise ValueError.

    """
    if seed is None or seed is np.random:
        return np.random.mtrand._rand
    if isinstance(seed, (numbers.Integral, np.integer)):
        return np.random.RandomState(seed)
    if isinstance(seed, np.random.RandomState):
        return seed
    raise ValueError('%r cannot be used to seed a numpy.random.RandomState'
                     ' instance' % seed)


def _get_n_jobs(n_jobs):
    """Get number of jobs for the computation.

    This function reimplements the logic of joblib to determine the actual
    number of jobs depending on the cpu count. If -1 all CPUs are used.
    If 1 is given, no parallel computing code is used at all, which is useful
    for debugging. For n_jobs below -1, (n_cpus + 1 + n_jobs) are used.
    Thus for n_jobs = -2, all CPUs but one are used.

    Parameters
    ----------
    n_jobs : int
        Number of jobs stated in joblib convention.

    Returns
    -------
    n_jobs : int
        The actual number of jobs as positive integer.

    """
    if n_jobs < 0:
        return max(cpu_count() + 1 + n_jobs, 1)
    elif n_jobs == 0:
        raise ValueError('Parameter n_jobs == 0 has no meaning.')
    else:
        return n_jobs


def _partition_estimators(n_estimators, n_jobs):
    """Private function used to partition estimators between jobs."""
    # Compute the number of jobs
    n_jobs = min(_get_n_jobs(n_jobs), n_estimators)

    # Partition estimators between jobs
    n_estimators_per_job = (n_estimators // n_jobs) * np.ones(n_jobs,
                                                              dtype=np.int)
    n_estimators_per_job[:n_estimators % n_jobs] += 1
    starts = np.cumsum(n_estimators_per_job)

    return n_jobs, n_estimators_per_job.tolist(), [0] + starts.tolist()


##
## TODO - predict
def _func_data_filter(X, apply_stack):
    terminals_list = []

    # function_name = apply_stack[-1][0].name
    for t in apply_stack[-1][1:]:
        if isinstance(t, int):
            terminals_list.append(X[:, t])
        elif isinstance(t, (float, np.float32, np.float64)):
            if t < 1.0:
                ###### ZBC: 这种情况一般是t是应该参数，例如ts_函数里面的天数d，其它基本没什么改动
                terminals_list.append(t)
            else:
                terminals_list.append(np.repeat(t, X.shape[0]))
        else:
            ###### ZBC: 这里是参入函数
            terminals_list.append(t)

    return terminals_list

def _func_data_filter_v2(X, apply_stack):
    terminals_list = []

    # function_name = apply_stack[-1][0].name
    for t in apply_stack[-1][1:]:
        if isinstance(t, str):
            terminals_list.append(X[t].values)
        elif isinstance(t, (float, np.float32, np.float64)):
            if t < 1.0:
                ###### ZBC: 这种情况一般是t是应该参数，例如ts_函数里面的天数d，其它基本没什么改动
                terminals_list.append(t)
            else:
                terminals_list.append(np.repeat(t, X.shape[0]))
        else:
            ###### ZBC: 这里是参入函数
            terminals_list.append(t)

    return terminals_list

def model_predict(X, program, keys=None, neu_keys=None, by_name=False, key_by_path=True):
    if len(program) > 1:
        base = program[1]
        gp_program = program[0]
    else:
        gp_program = program
        base = 5

    apply_stack = []
    for node in gp_program:
        if not isinstance(node, (str, int, float)):
            apply_stack.append([node])
        else:
            # Lazily evaluate later
            apply_stack[-1].append(node)

        while len(apply_stack[-1]) == apply_stack[-1][0].arity + 1:
            # Apply functions that have sufficient arguments
            function = apply_stack[-1][0]

            func_name = function.name

            if func_name.startswith('ts_'):
                used_func = eval('def_keyfunc_' + func_name)(keys, base=base, key_by_path=key_by_path)
            elif func_name.startswith('cs_'):
                if 'neu' in func_name:
                    used_func = eval('def_keyfunc_' + func_name)(keys, neu_keys, key_by_path=key_by_path)
                else:
                    used_func = eval('def_keyfunc_' + func_name)(keys, key_by_path=key_by_path)
            else:
                used_func = function

            ###### ZBC: 主要改动，见下面的函数
            if by_name:
                terminals = _func_data_filter_v2(X, apply_stack)
            else:
                terminals = _func_data_filter(X, apply_stack)
            # terminals = _func_data_filter(X, apply_stack)

            # TODO - run function
            intermediate_result = used_func(*terminals)

            if len(apply_stack) != 1:
                apply_stack.pop()
                apply_stack[-1].append(intermediate_result)
            else:
                return intermediate_result
                # print(intermediate_result)
                # break
